#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/epoll.h>
#include <fcntl.h>
#include <errno.h>

#define BUFFER_SIZE 1024    // 数据缓冲区大小
#define FILE_CHUNK_SIZE 512 // 文件分片大小
#define MAX_FILENAME 256    // 最大文件名长度
#define MAX_EVENTS 100      // epoll最大事件数
#define MAX_CLIENTS 1000    // 最大客户端数

#define ERR_LOG(errmsg) do { \
    perror(errmsg); \
    exit(1); \
} while(0)

// 定义消息类型（与客户端保持一致）
typedef enum {
    MSG_TEXT = 1,        // 普通文本消息
    MSG_FILE_SEND,       // 发送文件请求
    MSG_FILE_RECV,       // 接收文件请求
    MSG_FILE_DATA,       // 文件数据包
    MSG_FILE_END,        // 文件传输结束
    MSG_ACK,             // 确认消息
    MSG_ERROR            // 错误消息
} msg_type_t;

// 消息头结构体
typedef struct {
    msg_type_t type;           // 消息类型
    int seq_num;               // 序列号
    int data_len;              // 数据长度
    char filename[MAX_FILENAME]; // 文件名
} msg_header_t;

// 完整消息结构体
typedef struct {
    msg_header_t header;
    char data[FILE_CHUNK_SIZE];
} message_t;

// 客户端状态枚举
typedef enum {
    CLIENT_IDLE,           // 空闲状态
    CLIENT_RECEIVING_FILE, // 正在接收文件
    CLIENT_SENDING_FILE    // 正在发送文件
} client_state_t;

// 客户端信息结构体
typedef struct {
    int sockfd;
    struct sockaddr_in addr;
    client_state_t state;
    FILE *file_fp;                    // 文件指针
    char current_filename[MAX_FILENAME]; // 当前处理的文件名
    int expected_seq;                 // 期望的序列号
    int total_bytes;                  // 总字节数
} client_info_t;

// 全局客户端信息数组
client_info_t clients[MAX_CLIENTS];

// TCP发送函数 - 确保发送完整数据
ssize_t tcp_send_all(int sockfd, const void *buf, size_t len) {
    size_t total_sent = 0;
    const char *ptr = (const char *)buf;
    
    while (total_sent < len) {
        ssize_t sent = send(sockfd, ptr + total_sent, len - total_sent, 0);
        if (sent <= 0) {
            if (sent < 0) {
                perror("TCP发送失败");
            }
            return -1;
        }
        total_sent += sent;
    }
    return total_sent;
}

// TCP接收函数 - 确保接收完整数据
ssize_t tcp_recv_all(int sockfd, void *buf, size_t len) {
    size_t total_recv = 0;
    char *ptr = (char *)buf;
    
    while (total_recv < len) {
        ssize_t recvd = recv(sockfd, ptr + total_recv, len - total_recv, 0);
        if (recvd <= 0) {
            if (recvd < 0) {
                perror("TCP接收失败");
            }
            return -1;
        }
        total_recv += recvd;
    }
    return total_recv;
}

// 设置文件描述符为非阻塞模式
int set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) {
        perror("fcntl F_GETFL");
        return -1;
    }
    
    if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1) {
        perror("fcntl F_SETFL");
        return -1;
    }
    
    return 0;
}

// 初始化客户端信息
void init_client(client_info_t *client, int sockfd, struct sockaddr_in *addr) {
    client->sockfd = sockfd;
    client->addr = *addr;
    client->state = CLIENT_IDLE;
    client->file_fp = NULL;
    memset(client->current_filename, 0, MAX_FILENAME);
    client->expected_seq = 0;
    client->total_bytes = 0;
}

// 清理客户端资源
void cleanup_client(client_info_t *client) {
    if (client->file_fp) {
        fclose(client->file_fp);
        client->file_fp = NULL;
    }
    if (client->sockfd > 0) {
        close(client->sockfd);
        client->sockfd = -1;
    }
    client->state = CLIENT_IDLE;
    memset(client->current_filename, 0, MAX_FILENAME);
    client->expected_seq = 0;
    client->total_bytes = 0;
}

// 根据sockfd查找客户端
client_info_t* find_client(int sockfd) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].sockfd == sockfd) {
            return &clients[i];
        }
    }
    return NULL;
}

// 查找空闲的客户端slot
client_info_t* get_free_client_slot() {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].sockfd <= 0) {
            return &clients[i];
        }
    }
    return NULL;
}

// 发送文件给客户端
int send_file_to_client(client_info_t *client, const char *filename) {
    FILE *fp = fopen(filename, "rb");
    if (fp == NULL) {
        // 发送错误消息
        message_t error_msg;
        error_msg.header.type = MSG_ERROR;
        error_msg.header.seq_num = 0;
        error_msg.header.data_len = 0;
        strncpy(error_msg.header.filename, filename, MAX_FILENAME - 1);
        error_msg.header.filename[MAX_FILENAME - 1] = '\0';
        
        tcp_send_all(client->sockfd, &error_msg, sizeof(msg_header_t));
        
        printf("文件不存在或无法读取: %s (客户端 %s:%d)\n", 
               filename, inet_ntoa(client->addr.sin_addr), ntohs(client->addr.sin_port));
        return -1;
    }

    // 获取文件大小
    fseek(fp, 0, SEEK_END);
    long file_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    
    printf("开始发送文件给客户端 [%s:%d]: %s (大小: %ld 字节)\n",
           inet_ntoa(client->addr.sin_addr), ntohs(client->addr.sin_port), 
           filename, file_size);

    message_t msg;
    int seq_num = 0;
    int bytes_read;
    
    // 发送文件数据
    while ((bytes_read = fread(msg.data, 1, FILE_CHUNK_SIZE, fp)) > 0) {
        // 设置消息头
        msg.header.type = MSG_FILE_DATA;
        msg.header.seq_num = seq_num++;
        msg.header.data_len = bytes_read;
        strncpy(msg.header.filename, filename, MAX_FILENAME - 1);
        msg.header.filename[MAX_FILENAME - 1] = '\0';

        // 发送数据包
        if (tcp_send_all(client->sockfd, &msg, sizeof(msg_header_t) + bytes_read) == -1) {
            printf("发送文件数据失败\n");
            fclose(fp);
            return -1;
        }

        // 等待ACK确认
        message_t ack_msg;
        if (tcp_recv_all(client->sockfd, &ack_msg, sizeof(msg_header_t)) == -1) {
            printf("接收ACK失败\n");
            fclose(fp);
            return -1;
        }

        if (ack_msg.header.type != MSG_ACK || ack_msg.header.seq_num != seq_num - 1) {
            printf("收到错误的ACK\n");
            fclose(fp);
            return -1;
        }

        printf("已发送第 %d 个数据包给客户端 (%d 字节)\n", seq_num - 1, bytes_read);
    }

    // 发送文件结束标志
    msg.header.type = MSG_FILE_END;
    msg.header.seq_num = seq_num;
    msg.header.data_len = 0;
    strncpy(msg.header.filename, filename, MAX_FILENAME - 1);
    
    if (tcp_send_all(client->sockfd, &msg, sizeof(msg_header_t)) == -1) {
        printf("发送文件结束标志失败\n");
        fclose(fp);
        return -1;
    }

    // 等待最后的ACK
    message_t final_ack;
    if (tcp_recv_all(client->sockfd, &final_ack, sizeof(msg_header_t)) == -1) {
        printf("接收最终ACK失败\n");
    }

    fclose(fp);
    printf("文件发送完成给客户端: %s\n", filename);
    return 0;
}

// 处理客户端消息
void handle_client_message(client_info_t *client, message_t *msg) {
    message_t reply_msg, ack_msg;

    printf("\n[客户端 %s:%d] ", 
           inet_ntoa(client->addr.sin_addr), ntohs(client->addr.sin_port));

    switch (msg->header.type) {
        case MSG_TEXT:
            // 处理普通文本消息
            msg->data[msg->header.data_len] = '\0'; // 确保字符串结尾
            printf("收到文本消息: %s\n", msg->data);
            
            // 回复消息（添加服务器标识）
            char reply_text[BUFFER_SIZE];
            snprintf(reply_text, sizeof(reply_text), "服务器收到: %s", msg->data);
            
            reply_msg.header.type = MSG_TEXT;
            reply_msg.header.seq_num = 0;
            reply_msg.header.data_len = strlen(reply_text);
            strcpy(reply_msg.data, reply_text);
            
            if (tcp_send_all(client->sockfd, &reply_msg, 
                           sizeof(msg_header_t) + reply_msg.header.data_len) < 0) {
                printf("发送回复失败\n");
            }
            break;

        case MSG_FILE_RECV:
            // 客户端请求下载文件
            printf("请求下载文件: %s\n", msg->header.filename);
            send_file_to_client(client, msg->header.filename);
            break;

        case MSG_FILE_DATA:
            // 处理文件数据包
            if (client->state == CLIENT_IDLE) {
                // 第一个文件数据包，初始化接收状态
                char server_filename[MAX_FILENAME + 20];
                snprintf(server_filename, sizeof(server_filename), 
                        "server_%s", msg->header.filename);
                
                client->file_fp = fopen(server_filename, "wb");
                if (client->file_fp == NULL) {
                    printf("无法创建文件: %s\n", server_filename);
                    break;
                }
                
                strcpy(client->current_filename, server_filename);
                client->state = CLIENT_RECEIVING_FILE;
                client->expected_seq = 0;
                client->total_bytes = 0;
                
                printf("开始接收文件: %s -> %s\n", 
                       msg->header.filename, server_filename);
            }

            // 发送ACK
            ack_msg.header.type = MSG_ACK;
            ack_msg.header.seq_num = msg->header.seq_num;
            ack_msg.header.data_len = 0;
            
            if (tcp_send_all(client->sockfd, &ack_msg, sizeof(msg_header_t)) == -1) {
                printf("发送ACK失败\n");
            }

            // 检查序列号
            if (msg->header.seq_num != client->expected_seq) {
                printf("收到乱序的数据包，期望: %d，实际: %d\n",
                       client->expected_seq, msg->header.seq_num);
                break;
            }

            // 写入文件数据
            if (fwrite(msg->data, 1, msg->header.data_len, client->file_fp) != msg->header.data_len) {
                printf("写入文件失败\n");
                fclose(client->file_fp);
                client->file_fp = NULL;
                unlink(client->current_filename);
                client->state = CLIENT_IDLE;
                break;
            }

            client->total_bytes += msg->header.data_len;
            client->expected_seq++;
            printf("接收到第 %d 个数据包 (%d 字节)，总计: %d 字节\n",
                   msg->header.seq_num, msg->header.data_len, client->total_bytes);
            break;

        case MSG_FILE_END:
            // 文件传输结束
            if (client->state == CLIENT_RECEIVING_FILE && client->file_fp) {
                fclose(client->file_fp);
                client->file_fp = NULL;
                printf("文件接收完成: %s (总大小: %d 字节)\n", 
                       client->current_filename, client->total_bytes);
                client->state = CLIENT_IDLE;
            }
            
            // 发送最终ACK
            ack_msg.header.type = MSG_ACK;
            ack_msg.header.seq_num = msg->header.seq_num;
            ack_msg.header.data_len = 0;
            
            if (tcp_send_all(client->sockfd, &ack_msg, sizeof(msg_header_t)) == -1) {
                printf("发送最终ACK失败\n");
            }
            break;

        default:
            printf("收到未知类型的消息: %d\n", msg->header.type);
            break;
    }
}

int main(int argc, char const *argv[]) {
    if (argc < 3) {
        fprintf(stderr, "用法: %s <服务器IP> <端口>\n", argv[0]);
        exit(1);
    }

    // 初始化客户端数组
    for (int i = 0; i < MAX_CLIENTS; i++) {
        clients[i].sockfd = -1;
    }

    // 第一步：创建监听套接字
    int listen_fd;
    if ((listen_fd = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
        ERR_LOG("创建套接字失败");
    }

    // 设置套接字选项，允许地址重用
    int on = 1;
    if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) < 0) {
        ERR_LOG("设置套接字选项失败");
    }

    // 第二步：绑定套接字到服务器地址
    struct sockaddr_in serveraddr;
    serveraddr.sin_family = AF_INET;
    serveraddr.sin_addr.s_addr = inet_addr(argv[1]);
    serveraddr.sin_port = htons(atoi(argv[2]));

    if (bind(listen_fd, (struct sockaddr *)&serveraddr, sizeof(serveraddr)) == -1) {
        ERR_LOG("绑定失败");
    }

    // 第三步：开始监听
    if (listen(listen_fd, 10) == -1) {
        ERR_LOG("监听失败");
    }

    // 第四步：创建epoll实例
    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        ERR_LOG("创建epoll实例失败");
    }

    // 将监听套接字添加到epoll
    struct epoll_event ev, events[MAX_EVENTS];
    ev.events = EPOLLIN;
    ev.data.fd = listen_fd;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listen_fd, &ev) == -1) {
        ERR_LOG("添加监听套接字到epoll失败");
    }

    printf("TCP文件传输服务器启动成功\n");
    printf("监听地址: %s:%s\n", argv[1], argv[2]);
    printf("当前工作目录: %s\n", getcwd(NULL, 0));
    printf("等待客户端连接...\n\n");

    // 主服务循环
    while (1) {
        int nfds = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
        if (nfds == -1) {
            perror("epoll_wait失败");
            continue;
        }

        for (int i = 0; i < nfds; i++) {
            int fd = events[i].data.fd;

            if (fd == listen_fd) {
                // 新的客户端连接
                struct sockaddr_in client_addr;
                socklen_t client_len = sizeof(client_addr);
                int client_fd = accept(listen_fd, (struct sockaddr *)&client_addr, &client_len);
                
                if (client_fd == -1) {
                    perror("接受客户端连接失败");
                    continue;
                }

                printf("新客户端连接: %s:%d (fd=%d)\n", 
                       inet_ntoa(client_addr.sin_addr), 
                       ntohs(client_addr.sin_port), client_fd);

                // 查找空闲的客户端slot
                client_info_t *client = get_free_client_slot();
                if (client == NULL) {
                    printf("服务器已达到最大客户端数量限制\n");
                    close(client_fd);
                    continue;
                }

                // 初始化客户端信息
                init_client(client, client_fd, &client_addr);

                // 设置非阻塞模式（可选）
                // set_nonblocking(client_fd);

                // 将客户端套接字添加到epoll
                ev.events = EPOLLIN;
                ev.data.fd = client_fd;
                if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_fd, &ev) == -1) {
                    perror("添加客户端套接字到epoll失败");
                    cleanup_client(client);
                    continue;
                }

            } else {
                // 客户端数据到达
                client_info_t *client = find_client(fd);
                if (client == NULL) {
                    printf("未找到客户端信息 (fd=%d)\n", fd);
                    continue;
                }

                message_t msg;
                
                // 先接收消息头
                ssize_t header_len = tcp_recv_all(fd, &msg.header, sizeof(msg_header_t));
                if (header_len <= 0) {
                    // 客户端断开连接
                    printf("客户端 [%s:%d] 断开连接\n", 
                           inet_ntoa(client->addr.sin_addr), 
                           ntohs(client->addr.sin_port));
                    
                    // 从epoll中移除
                    epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, NULL);
                    
                    // 清理客户端资源
                    cleanup_client(client);
                    continue;
                }

                // 如果有数据部分，接收数据
                if (msg.header.data_len > 0) {
                    if (tcp_recv_all(fd, msg.data, msg.header.data_len) <= 0) {
                        printf("接收客户端数据失败\n");
                        
                        // 从epoll中移除
                        epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, NULL);
                        
                        // 清理客户端资源
                        cleanup_client(client);
                        continue;
                    }
                }

                // 处理客户端消息
                handle_client_message(client, &msg);
            }
        }
    }

    // 清理资源
    close(epoll_fd);
    close(listen_fd);
    
    // 清理所有客户端
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].sockfd > 0) {
            cleanup_client(&clients[i]);
        }
    }

    return 0;
}